package tinyiotserver

import (
	"bufio"
	"bytes"
	"context"
	"crypto/rand"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"runtime"
	"sync"
	"time"

	"gitee.com/mqyqingkong/flowprocess"
	"gitee.com/mqyqingkong/tiny-iot-server/common"
)

type Server interface {
	HandlerRegistrar
	Start() error
	// Submit 在Server启动后，直接提交消息
	Submit(msg *common.Msg) (*common.Msg, error)
	Close() error
}

type ServerOption func(s *server)

// WithTopicQueueSize 配置topicQueueSize
func WithTopicQueueSize(topicQueueSize map[string]int) ServerOption {
	return func(s *server) {
		s.topicQueueSize = topicQueueSize
	}
}

// WithDefaultQueueSize 配置默认的topicQueueSize
func WithDefaultQueueSize(defaultQueueSize int) ServerOption {
	return func(s *server) {
		s.defaultQueueSize = defaultQueueSize
	}
}

// WithSocketWriteDeadline 配置端口超时
func WithSocketWriteDeadline(socketWriteDeadline time.Duration) ServerOption {
	return func(s *server) {
		s.socketWriteDeadline = socketWriteDeadline
	}
}

// WithHeartbeatDuration 心跳周期
func WithHeartbeatDuration(heartbeatDuration time.Duration) ServerOption {
	return func(s *server) {
		if heartbeatDuration >= time.Second {
			s.heartbeatDuration = heartbeatDuration
		}
	}
}

// WithNetwork 必须是 "tcp", "tcp4", "tcp6", "unix",  "unixpacket".
func WithNetwork(network string) ServerOption {
	return func(s *server) {
		s.network = network
	}
}

type MsgHandlerSelector interface {
	// DefineProcessorInTaskChan 为每个processor定义一个输入chan，返回值为输入chan的id
	DefineProcessorInTaskChan() (processorInTaskChanIndex []int)
	// SelectInTaskChanIndex 将task分配给某个chan。
	SelectInTaskChanIndex(topic string, payload []byte) (inTaskChanIndex int)
}

type msgHandlerProcessorSelectorAdapter struct {
	msgHandlerSelector MsgHandlerSelector
}

func (s *msgHandlerProcessorSelectorAdapter) DefineProcessorInTaskChan() (processorInTaskChanIndex []int) {
	return s.msgHandlerSelector.DefineProcessorInTaskChan()
}

func (s *msgHandlerProcessorSelectorAdapter) SelectInTaskChanIndex(task flowprocess.Task) (inTaskChanIndex int) {
	msg := task.(*msgConn).message
	return s.msgHandlerSelector.SelectInTaskChanIndex(string(msg.Body.Topic), msg.Body.Payload)
}

func WithMsgHandlerSelector(msgHandlerSelector MsgHandlerSelector) ServerOption {
	return func(s *server) {
		s.msgHandlerNum = len(msgHandlerSelector.DefineProcessorInTaskChan())
		s.msgHandlerSelector = &msgHandlerProcessorSelectorAdapter{
			msgHandlerSelector: msgHandlerSelector,
		}
	}
}

func WithTLS(keyPath, certPath string) ServerOption {
	return func(s *server) {
		s.tlsKeyPath = keyPath
		s.tlsCertPath = certPath
	}
}

func NewServer(address string, options ...ServerOption) Server {
	ctx, cancelFunc := context.WithCancel(context.Background())
	s := &server{
		address:             address,
		ctx:                 ctx,
		cancel:              cancelFunc,
		flow:                flowprocess.NewFlow(),
		msgHandlers:         map[string][]MessageHandler{},
		socketWriteDeadline: 5 * time.Second,
		topicQueueSize:      map[string]int{},
		defaultQueueSize:    512,
		topics:              map[string]Queue{},
		heartbeatDuration:   5 * time.Second,
		connClosed:          map[net.Conn]chan struct{}{},
		network:             "tcp",
	}
	for i := 0; i < len(options); i++ {
		options[i](s)
	}
	return s
}

type server struct {
	address                 string
	ctx                     context.Context
	cancel                  context.CancelFunc
	flow                    flowprocess.Flow //流处理框架，参考 gitee.com/mqyqingkong/flowprocess
	msgHandlers             map[string][]MessageHandler
	socketWriteDeadline     time.Duration
	topicQueueSize          map[string]int
	defaultQueueSize        int
	topics                  map[string]Queue
	heartbeatDuration       time.Duration
	connClosed              map[net.Conn]chan struct{} //监听connection是否已关闭
	lock                    sync.RWMutex
	network                 string
	processNode             flowprocess.Node
	msgHandlerNum           int
	msgHandlerSelector      flowprocess.ProcessorSelector
	tlsKeyPath, tlsCertPath string
}

func (s *server) listen() (net.Listener, error) {

	if len(s.tlsKeyPath) == 0 || len(s.tlsCertPath) == 0 {
		return net.Listen(s.network, s.address)
	}

	crt, err := tls.LoadX509KeyPair(s.tlsCertPath, s.tlsKeyPath)
	if err != nil {
		return nil, err
	}
	tlsConfig := &tls.Config{}
	tlsConfig.Certificates = []tls.Certificate{crt}
	tlsConfig.Time = time.Now
	tlsConfig.Rand = rand.Reader
	return tls.Listen(s.network, s.address, tlsConfig)
}

func (s *server) Start() error {
	listener, err := s.listen()
	if err != nil {
		return err
	}
	//Node0,监听端口（1个go-routine）
	netListenHandler := &NetListenHandler{
		ctx:      s.ctx,
		listener: listener,
	}
	s.flow.AddNodeTaskHandlers(1000, netListenHandler)

	//Node1,组装消息（2个go-routine）
	messageHandler := &ReadMessageHandler{
		server: s,
	}
	s.flow.AddNodeTaskHandlers(1000, messageHandler, messageHandler)

	//Node2 处理和回复消息（NumCPU个go-routine）
	processMessageHandler := &ProcessMessageHandler{
		server: s,
	}
	msgHandlerNum := s.msgHandlerNum
	if msgHandlerNum <= 0 {
		msgHandlerNum = runtime.NumCPU()
	}
	processHandlers := make([]flowprocess.TaskHandler, msgHandlerNum)
	for i := 0; i < msgHandlerNum; i++ {
		processHandlers[i] = processMessageHandler
	}
	s.processNode = s.flow.AddNodeTaskHandlers(1000, processHandlers...)
	s.processNode.SetProcessorSelector(s.msgHandlerSelector)

	//Node3 写入消息队列、读取消息队列（1个go-routine，不存在并发访问内存问题）
	queueMessageHandler := NewQueueMessageHandler(s)
	s.flow.AddNodeTaskHandlers(0, queueMessageHandler)

	s.flow.Start()

	log.Println("server started")

	s.flow.Await()
	return nil
}

func (s *server) Submit(msg *common.Msg) (*common.Msg, error) {
	now := time.Now()
	expiredTime := time.Time{}
	if msg.Header.ExpireAfterNow != 0 {
		expiredTime = now.Add(time.Duration(msg.Header.ExpireAfterNow) * time.Millisecond)
	}

	done := make(chan struct{})
	var ret *common.Msg

	err := s.processNode.SubmitTask(&msgConn{
		message: &Message{
			Msg:          msg,
			ReceivedTime: now,
			ExpiredTime:  expiredTime,
		},
		reply: func(replyMsg *Message) error {
			ret = replyMsg.Msg
			close(done)
			return nil
		},
	})
	if err != nil {
		close(done)
		return nil, err
	}
	needReplyToClient := msg.Header.Type != common.POST_NO_RESPONSE
	if needReplyToClient {
		select {
		case <-s.ctx.Done():
			return ret, nil
		case <-done:
			return ret, nil
		case <-time.After(10 * time.Second):
			return nil, errors.New("time out waiting for result")
		}
	} else {
		close(done)
		return nil, nil
	}
}

func (s *server) Close() error {
	log.Println("server begins destroying")
	s.cancel()
	if s.network == "unix" {
		os.Remove(s.address)
	}
	return nil
}

func (s *server) RegisterGlobalPreHandler(handlers ...MessageHandler) {
	path := "GlobalPreHandler"
	s.doRegisterHandler(path, handlers...)
}

func (s *server) RegisterHandler(topic string, handlers ...MessageHandler) {
	s.doRegisterHandler(topic, handlers...)
}

func (s *server) doRegisterHandler(path string, handlers ...MessageHandler) {
	if msgHs, ok := s.msgHandlers[path]; ok {
		msgHs = append(msgHs, handlers...)
		s.msgHandlers[path] = msgHs
	} else {
		s.msgHandlers[path] = handlers
	}
}

func (s *server) GetHandlers(topic string) []MessageHandler {
	var hs []MessageHandler
	//先获取通用的handler
	path := "GlobalPreHandler"
	hs = append(hs, s.msgHandlers[path]...)

	//再获取特定的handler
	hs = append(hs, s.msgHandlers[topic]...)

	return hs
}

type msgConn struct {
	conn    net.Conn
	message *Message
	// reply 向客户端写入响应消息
	reply func(replyMsg *Message) error
}

// replyToClientAsyn 向客户端写入响应消息
func (s *server) replyToClientAsyn(responseMsg *Message, mc *msgConn) {
	go func() {
		replyToClientMsg := s.onResponseCallback(responseMsg, mc.message)
		s.replyToClient(replyToClientMsg, mc)
	}()
}

// replyToClientAsyn 向客户端写入响应消息
func (s *server) replyToClientDirectlyAsyn(responseMsg *Message, mc *msgConn) {
	go func() {
		s.replyToClient(responseMsg, mc)
	}()
}

func (s *server) onResponseCallback(responseMsg *Message, reqMsg *Message) (replyToClientMsg *Message) {
	topic := string(reqMsg.Body.Topic)
	//处理并回复消息
	handlers := s.GetHandlers(topic)
	replyToClientMsg = responseMsg
	for i := len(handlers) - 1; i >= 0; i-- {
		replyToClientMsg = handlers[i].OnResponse(replyToClientMsg, reqMsg)
	}
	return
}

// replyToClient 向客户端写入响应消息
func (s *server) replyToClient(replyMsg *Message, mc *msgConn) error {
	if replyMsg == nil {
		return nil
	}
	if mc.reply != nil {
		replyMsg.Header.Type = common.RESPONSE
		err := mc.reply(replyMsg)
		if err != nil {
			log.Printf("[ERR] reply to client error: %v, msg: %s \n", err, replyMsg)
		}
		return err
	}
	return nil
}

func (s *server) getTopicQueueSize(topic string) int {
	if size, found := s.topicQueueSize[topic]; found {
		return size
	}
	return s.defaultQueueSize
}

// LoadCreateQueue 查询topic对应的队列，找不到则创建一个
func (s *server) LoadCreateQueue(topic string, msgType common.MsgType) Queue {
	queue, ok := s.topics[topic]
	isQueryStream := msgType.IsQueryStream()
	if ok {
		// 需要加锁queue
		if isQueryStream {
			if _, synced := queue.(*syncQueue); !synced {
				queue = NewSyncBlockingQueue(queue)
				s.topics[topic] = queue
			}
		}
	} else {
		queueMaxLen := s.getTopicQueueSize(topic)
		queue = NewOrderQueue(queueMaxLen)
		// queue = NewListQueue(queueMaxLen)
		// 需要加锁queue
		if isQueryStream {
			queue = NewSyncBlockingQueue(queue)
		}
		s.topics[topic] = queue
	}
	return queue
}

// addConnection 当有新的conn接入时，保存其关闭状态，以便在后续推送数据时，及时退出goroutine
func (s *server) addConnection(conn net.Conn) {
	s.lock.Lock()
	defer s.lock.Unlock()
	s.connClosed[conn] = make(chan struct{})
}

func (s *server) removeConnection(conn net.Conn) {
	s.lock.Lock()
	defer s.lock.Unlock()
	connC := s.connClosed[conn]
	if connC != nil {
		close(connC)
		delete(s.connClosed, conn)
	}
}

// connectionClosed 监听conn是否已关闭
func (s *server) connectionClosed(conn net.Conn) <-chan struct{} {
	s.lock.RLock()
	defer s.lock.RUnlock()
	connC := s.connClosed[conn]
	if connC == nil {
		connC = make(chan struct{})
		close(connC)
	}
	return connC
}

type Message struct {
	*common.Msg
	ReceivedTime time.Time
	ExpiredTime  time.Time
	No           int64
}

func (m *Message) Copy() *Message {
	ret := &Message{
		Msg:          m.Msg.Copy(),
		ReceivedTime: m.ReceivedTime,
		ExpiredTime:  m.ExpiredTime,
		No:           m.No,
	}
	return ret
}

func (m *Message) SetTopic(topic string) {
	m.Body.Topic = []byte(topic)
	m.InitLength()
}

func (m *Message) SetPayload(payload []byte) {
	m.Body.Payload = payload
	m.InitLength()
}

func (m *Message) IsExpired() bool {
	if m.ExpiredTime.IsZero() {
		return false
	}
	return m.ExpiredTime.Before(time.Now())
}

func (m *Message) IsEmpty() bool {
	return bytes.Equal(m.Body.Payload, []byte(common.EMPTY))
}

func (msg *Message) String() string {
	end := 10
	if int(msg.Body.PayloadLen) < end {
		end = int(msg.Body.PayloadLen)
	}
	if msg.ExpiredTime.IsZero() {
		return fmt.Sprintf("[Version:%d, ReqType:%s, ClientId:%d, ReceivedTime:%s, Topic: %s, Payload: %s...]",
			msg.Header.Version, msg.Header.Type, msg.Header.ClientId, msg.ReceivedTime.Format("2006-01-02 15:04:05"), string(msg.Body.Topic), string(msg.Body.Payload[:end]))
	}
	return fmt.Sprintf("[Version:%d, ReqType:%s, ClientId:%d, ReceivedTime:%s, ExpiredTime:%s, Topic: %s, Payload: %s...]",
		msg.Header.Version, msg.Header.Type, msg.Header.ClientId, msg.ReceivedTime.Format("2006-01-02 15:04:05"), msg.ExpiredTime.Format("2006-01-02 15:04:05"), string(msg.Body.Topic), string(msg.Body.Payload[:end]))
}

func SuccessMsg(topic string) *Message {
	ret := &Message{
		Msg: common.SuccessMsg(topic),
	}
	return ret
}

func FailMsg(topic string, errorInfo string) *Message {
	ret := &Message{
		Msg: common.FailMsg(topic, errorInfo),
	}
	return ret
}

func ResponseMsg(topic string, payload []byte) *Message {
	ret := &Message{
		Msg: common.ResponseMsg(topic, payload),
	}
	return ret
}

// Node0,监听端口
type NetListenHandler struct {
	flowprocess.TaskHandlerAdapter
	ctx      context.Context
	listener net.Listener
}

func (h *NetListenHandler) Handle(inTask flowprocess.Task, dispatch func(outTask flowprocess.Task) error) (err error) {
	for {
		select {
		case <-h.ctx.Done():
			h.listener.Close()
			return errors.New("task canceled by user")
		default:
			conn, err := h.listener.Accept()
			if err != nil {
				log.Print("[ERR] accept failed, err:", err)
				continue
			}
			//转发给下一个Node
			dispatch(conn)
		}
	}
}

// Node1,组装消息
type ReadMessageHandler struct {
	flowprocess.TaskHandlerAdapter
	server *server
}

func (h *ReadMessageHandler) Handle(inTask flowprocess.Task, dispatch func(outTask flowprocess.Task) error) (err error) {
	conn := inTask.(net.Conn)

	go func(conn net.Conn) {
		h.server.addConnection(conn)
		defer h.server.removeConnection(conn)

		reply := func(replyMsg *Message) error {
			conn.SetWriteDeadline(time.Now().Add(h.server.socketWriteDeadline))
			_, err := conn.Write(replyMsg.Bytes())
			return err
		}

		r := bufio.NewReader(conn)
		var msgNo int64 = 0
		for {
			msg, err := common.ReadMsg(r)
			if err == io.EOF {
				return
			}
			if err != nil {
				log.Printf("[ERR] read message error: %v", err)
				return
			}
			now := time.Now()
			expiredTime := time.Time{}
			if msg.Header.ExpireAfterNow != 0 {
				expiredTime = now.Add(time.Duration(msg.Header.ExpireAfterNow) * time.Millisecond)
			}
			message := &Message{
				Msg:          msg,
				ReceivedTime: now,
				ExpiredTime:  expiredTime,
				No:           msgNo,
			}
			msgNo++
			//转发给下一个Node
			dispatch(&msgConn{
				conn:    conn,
				message: message,
				reply:   reply,
			})
		}
	}(conn)

	return nil
}

// Node2 处理和回复消息
type ProcessMessageHandler struct {
	flowprocess.TaskHandlerAdapter
	server *server
}

func (h *ProcessMessageHandler) Handle(inTask flowprocess.Task, dispatch func(outTask flowprocess.Task) error) error {
	mc := inTask.(*msgConn)
	reqType := mc.message.Header.Type
	topic := string(mc.message.Body.Topic)
	//处理并回复消息
	handlers := h.server.GetHandlers(topic)
	if len(handlers) > 0 {
		for i := 0; i < len(handlers); i++ {
			newMsg := handlers[i].OnReceive(mc.message)
			if newMsg == nil {
				continue
			}
			//此消息是RESPONSE类型，终止后续消息执行，直接回复给客户端
			if newMsg.Header.Type == common.RESPONSE {
				if reqType == common.POST_WAIT_RESPONSE || reqType.IsQuery() {
					h.server.replyToClientDirectlyAsyn(newMsg, mc)
				}
				return nil
			}
			//转发给下一个Node
			dispatch(&msgConn{
				conn:    mc.conn,
				message: newMsg,
				reply:   mc.reply,
			})
		}
	} else {
		//转发给下一个Node
		dispatch(mc)
	}
	if reqType == common.POST_WAIT_RESPONSE {
		newMsg := SuccessMsg(topic)
		h.server.replyToClientDirectlyAsyn(newMsg, mc)
	}
	return nil
}

func NewQueueMessageHandler(server *server) *QueueMessageHandler {
	h := &QueueMessageHandler{
		server:                server,
		topicQueueBroadcastor: map[string]*QueueBroadcastor{},
		broadcastToTarget:     sync.Map{},
	}
	return h
}

// Node3 将消息放入队列, 单goroutine运行，一般不用加锁
type QueueMessageHandler struct {
	flowprocess.TaskHandlerAdapter
	server                *server
	topicQueueBroadcastor map[string]*QueueBroadcastor // topic:*QueueBroadcastor
	broadcastToTarget     sync.Map                     // 同一个connection，可以请求多次common.QUERY_STREAM_GROUP，此时会更改广播的目的地
}

func (h *QueueMessageHandler) Handle(inTask flowprocess.Task, dispatch func(outTask flowprocess.Task) error) error {
	mc := inTask.(*msgConn)
	reqMsg := mc.message
	reqType := reqMsg.Header.Type
	topic := string(reqMsg.Body.Topic)
	q := h.server.LoadCreateQueue(topic, reqType)

	if reqType == common.QUERY {
		//查询消息
		ret := h.queryMsg(q, reqMsg)
		h.server.replyToClientAsyn(ret, mc)
	} else if reqType == common.QUERY_POP {
		//查询消息
		ret := h.queryPopMsg(q, reqMsg)
		h.server.replyToClientAsyn(ret, mc)
	} else if reqType.IsPost() {
		//写入消息
		q.PushFront(reqMsg)
	} else if reqType == common.QUERY_STREAM {
		h.processQueryStream(q, mc)
	} else if reqType == common.QUERY_STREAM_GROUP {
		h.processQueryStreamGroup(q, mc)
	}
	return nil
}

func (h *QueueMessageHandler) loadCreateQueueBroadcastor(topic string, queue Queue) *QueueBroadcastor {
	broadcastor := h.topicQueueBroadcastor[topic]
	if broadcastor == nil {
		broadcastor = NewQueueBroadcastor(h.server.ctx, queue, topic, h.server.heartbeatDuration)
		h.topicQueueBroadcastor[topic] = broadcastor
	}
	return broadcastor
}

type broadcastTargetInfo struct {
	broadcastToTarget func(msg *Message) error
	listenTopic       string
}

func (h *QueueMessageHandler) processQueryStreamGroup(queue Queue, mc *msgConn) {
	topic := string(mc.message.Body.Topic)

	broadcastToNewTarget := func(msg *Message) error {
		replyToClientMsg := h.server.onResponseCallback(msg, mc.message)
		err := h.server.replyToClient(replyToClientMsg, mc)
		return err
	}
	broadcastToNewTargetInfo := &broadcastTargetInfo{
		broadcastToTarget: broadcastToNewTarget,
		listenTopic:       topic,
	}

	// 之前是否有广播任务
	broadcastToTargetInfo, hasBroadTaskBefore := h.broadcastToTarget.Load(mc.conn)
	h.broadcastToTarget.Store(mc.conn, broadcastToNewTargetInfo) //新增/更新广播target
	if hasBroadTaskBefore && broadcastToTargetInfo.(*broadcastTargetInfo).listenTopic == topic {
		// 如果之前有广播任务，且监听的topic没变化，则不必再广播，只修改广播target
		log.Printf("[ClientId:%d]:Topic[%s] has broadcast task before. Broadcast target has been updated.", mc.message.Header.ClientId, topic)
		return
	}

	broadcastor := h.loadCreateQueueBroadcastor(topic, queue)

	broadcastToTarget := func(msg *Message) error {
		broadcastToNewTargetInfo, found := h.broadcastToTarget.Load(mc.conn)
		if found {
			broadcastToNewTarget := broadcastToNewTargetInfo.(*broadcastTargetInfo)
			if broadcastToNewTarget.listenTopic != topic {
				return nil
			}
			return broadcastToNewTarget.broadcastToTarget(msg)
		}
		return nil
	}
	stopCondition := func() <-chan struct{} {
		broadcastToNewTargetInfo, found := h.broadcastToTarget.Load(mc.conn)
		if found { //如果conn监听的topic变化，则停止对该conn广播
			listenTopic := broadcastToNewTargetInfo.(*broadcastTargetInfo).listenTopic
			if listenTopic != topic {
				log.Printf("[ClientId:%d]:Topic[%s-->%s] switched when broadcast task is running. Topic[%s] listening is finished.", mc.message.Header.ClientId, topic, listenTopic, topic)
				done := make(chan struct{})
				close(done)
				return done
			}
		}

		return h.server.connectionClosed(mc.conn)
	}

	broadcastor.Broadcast(broadcastToTarget, stopCondition)
}

func (h *QueueMessageHandler) processQueryStream(topicQueue Queue, mc *msgConn) {
	go func() {
		heartbeatTimer := time.NewTimer(h.server.heartbeatDuration)
		defer func() {
			heartbeatTimer.Stop()
			log.Printf("processQueryStream finished for topic: %s \n", mc.message.Body.Topic)
		}()

		for {
			select {
			case <-h.server.ctx.Done():
				return
			case <-h.server.connectionClosed(mc.conn):
				return
			case <-topicQueue.Wait():
				replyMsg, found := topicQueue.PopBack()
				if found {
					if !heartbeatTimer.Stop() {
						<-heartbeatTimer.C
					}
					replyToClientMsg := h.server.onResponseCallback(replyMsg, mc.message)
					err := h.server.replyToClient(replyToClientMsg, mc)
					if err != nil {
						topicQueue.PushBack(replyMsg)
						return
					}
					heartbeatTimer.Reset(h.server.heartbeatDuration)
				}
			case <-heartbeatTimer.C:
				replyMsg := ResponseMsg(string(mc.message.Body.Topic), []byte(common.HEARTBEAT))
				log.Printf("heartbeat: %s", replyMsg)
				err := h.server.replyToClient(replyMsg, mc)
				if err != nil {
					return
				}
				heartbeatTimer.Reset(h.server.heartbeatDuration)
			}
		}

	}()
}

func (h *QueueMessageHandler) queryMsg(q Queue, reqMsg *Message) *Message {
	retMsg, found := q.PopBack()
	for found && retMsg.IsExpired() {
		//如果消息过期，丢弃消息
		retMsg, found = q.PopBack()
	}

	if found {
		//如果是QUERY，则放回消息。
		q.PushBack(retMsg)
		return retMsg
	} else {
		return ResponseMsg(string(reqMsg.Body.Topic), []byte(common.EMPTY))
	}
}

func (h *QueueMessageHandler) queryPopMsg(q Queue, reqMsg *Message) *Message {
	retMsg, found := q.PopBack()
	for found && retMsg.IsExpired() {
		//如果消息过期，丢弃消息
		retMsg, found = q.PopBack()
	}

	if found {
		return retMsg
	} else {
		return ResponseMsg(string(reqMsg.Body.Topic), []byte(common.EMPTY))
	}
}
